import base64
import re
import io
import pytesseract
from PIL import Image, ImageFilter, ImageEnhance
import numpy as np


# 设置Tesseract路径（根据您的环境配置）
# pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'

def base64_to_image(base64_str):
    """将包含头部的Base64字符串转换为PIL图像对象"""
    # 检查并移除头部信息（如果存在）
    if any(base64_str.startswith(prefix) for prefix in [
        'data:image/png',
        'data:image/jpeg',
        'data:image/jpg',
        'data:image/gif'
    ]):
        base64_str = re.sub(r'^data:image/[^;]+;base64,', '', base64_str)

    # 解码Base64
    try:
        image_data = base64.b64decode(base64_str)
        return Image.open(io.BytesIO(image_data))
    except Exception as e:
        raise ValueError(f"Base64解码失败: {str(e)}")


def preprocess_captcha_image(img):
    """针对验证码优化的预处理"""
    # 转换为灰度图
    if img.mode != 'L':
        img = img.convert('L')

    # 增强对比度
    enhancer = ImageEnhance.Contrast(img)
    img = enhancer.enhance(2.0)  # 200%对比度增强

    # 二值化处理
    img_np = np.array(img)
    threshold = np.mean(img_np) * 0.8  # 动态计算阈值
    img = img.point(lambda px: 255 if px > threshold else 0)

    # 去噪
    img = img.filter(ImageFilter.MedianFilter(size=3))

    return img


def recognize_captcha(base64_str):
    """
    识别Base64编码的验证码图像
    :param base64_str: Base64编码的字符串（带或不带头部）
    :return: 识别出的验证码文本
    """
    try:
        # 1. 将Base64转换为PIL图像
        img = base64_to_image(base64_str)

        # 2. 图像预处理
        processed_img = preprocess_captcha_image(img)

        # 3. 使用Tesseract进行OCR识别
        # 尝试多种配置以提高识别率
        configs = [
            '--psm 7 --oem 3',  # 单行文本
            '--psm 8 --oem 3',  # 单个单词
            '--psm 10 --oem 3'  # 单字符
        ]

        best_result = ""
        for config in configs:
            try:
                text = pytesseract.image_to_string(
                    processed_img,
                    config=config
                )

                # 清理结果：移除空格和特殊字符，只保留字母数字
                cleaned_text = ''.join(c for c in text if c.isalnum()).upper()

                # 如果结果合理则使用
                if cleaned_text and 4 <= len(cleaned_text) <= 8:  # 验证码通常4-8个字符
                    return cleaned_text

            except Exception:
                continue

        # 如果所有配置都失败，使用默认配置
        text = pytesseract.image_to_string(processed_img)
        cleaned_text = ''.join(c for c in text if c.isalnum()).upper()
        return cleaned_text if cleaned_text else "无法识别"

    except Exception as e:
        return f"识别过程中出错: {str(e)}"


def get_truncated_base64(base64_str, length=10):
    """截断Base64字符串，只显示部分内容"""
    if len(base64_str) > length:
        return base64_str[:length] + "..."
    return base64_str



if __name__ == '__main__':
    # 您的 Base64 编码的验证码
    base64_captcha = ''
    # 调用函数识别验证码
    result = recognize_captcha(base64_captcha)

    # 打印识别结果
    print("识别出的验证码是:", result)